package com.jxpanda.spring.module.auth.config;

import com.jxpanda.spring.module.auth.config.properties.JwtProperties;
import com.jxpanda.spring.module.auth.config.properties.TokenProperties;
import com.jxpanda.spring.module.auth.core.token.JwtReactiveOAuth2TokenGenerator;
import com.jxpanda.spring.module.auth.core.token.ReactiveJwtCustomizer;
import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jose.jwk.JWKMatcher;
import com.nimbusds.jose.jwk.JWKSelector;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.proc.SecurityContext;
import com.nimbusds.jwt.SignedJWT;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.autoconfigure.AutoConfigureBefore;
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnExpression;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import org.springframework.security.oauth2.jwt.NimbusJwtEncoder;
import org.springframework.security.oauth2.jwt.NimbusReactiveJwtDecoder;
import org.springframework.security.oauth2.jwt.ReactiveJwtDecoder;
import reactor.core.publisher.Flux;

import java.util.function.Function;

@RequiredArgsConstructor
@AutoConfigureBefore(TokenConfigure.class)
@ConditionalOnProperty(prefix = "panda.spring.auth.token.jwt", name = "enable", havingValue = "true")
public class JwtConfigure {

    private final TokenProperties tokenProperties;

    private final JwtProperties jwtProperties;


    @Bean("jwkSource")
    @ConditionalOnMissingBean
    public JWKSource<SecurityContext> jwkSource() {
        return jwtProperties.loadJwkSource();
    }

    @Bean("jwtEncoder")
    @ConditionalOnBean(name = "jwkSource")
    @ConditionalOnMissingBean(name = "jwtEncoder")
    public JwtEncoder jwtEncoder(@Qualifier("jwkSource") JWKSource<SecurityContext> jwkSource) {
        return new NimbusJwtEncoder(jwkSource);
    }

    @Bean("reactiveJwtDecoder")
    @ConditionalOnBean(name = "jwkSource")
    @ConditionalOnMissingBean(name = "reactiveJwtDecoder")
    public ReactiveJwtDecoder reactiveJwtDecoder(@Qualifier("jwkSource") JWKSource<SecurityContext> jwkSource) {
        Function<SignedJWT, Flux<JWK>> jwkSourceFunction = signedJWT -> {
            try {
                // 从 SignedJWT 的头部创建 JWKMatcher
                JWKMatcher jwkMatcher = JWKMatcher.forJWSHeader(signedJWT.getHeader());
                JWKSelector jwkSelector = new JWKSelector(jwkMatcher);
                // 使用 JWKSource 获取匹配的 JWK 列表
                // 返回 Flux<JWK>
                return Flux.fromIterable(jwkSource.get(jwkSelector, null));
            } catch (Exception e) {
                return Flux.error(new RuntimeException("Failed to retrieve JWKs", e));
            }
        };
        return NimbusReactiveJwtDecoder.withJwkSource(jwkSourceFunction)
                .jwsAlgorithm(jwtProperties.getAlgorithm().getJwsAlgorithm())
                .build();
    }

    @Bean("accessTokenGenerator")
    @ConditionalOnBean(name = "jwtEncoder")
    @ConditionalOnMissingBean(name = "accessTokenGenerator")
    public JwtReactiveOAuth2TokenGenerator accessTokenGenerator(
            @Qualifier("jwtEncoder") JwtEncoder jwtEncoder,
            @Qualifier("reactiveJwtCustomizer")
            @Autowired(required = false) ReactiveJwtCustomizer reactiveJwtCustomizer) {
        return JwtReactiveOAuth2TokenGenerator.builder()
                .expiresIn(tokenProperties.getAccess().getExpiresIn())
                .jwtEncoder(jwtEncoder)
                .jwtProperties(jwtProperties)
                .reactiveJwtCustomizer(reactiveJwtCustomizer)
                .build();
    }

    @Bean("refreshTokenGenerator")
    @ConditionalOnBean(name = "jwtEncoder")
    @ConditionalOnMissingBean(name = "refreshTokenGenerator")
    @ConditionalOnExpression("${panda.spring.auth.token.refresh.enable:true}")
    @ConditionalOnProperty(prefix = "panda.spring.auth.token.refresh", name = "mode", havingValue = "JWT")
    public JwtReactiveOAuth2TokenGenerator refreshTokenGenerator(@Qualifier("jwtEncoder") JwtEncoder jwtEncoder) {
        return JwtReactiveOAuth2TokenGenerator.builder()
                .expiresIn(tokenProperties.getRefresh().getExpiresIn())
                .jwtEncoder(jwtEncoder)
                .jwtProperties(jwtProperties)
                .build();
    }

    @Bean("reactiveJwtCustomizer")
    @ConditionalOnMissingBean(name = "reactiveJwtCustomizer")
    public ReactiveJwtCustomizer reactiveJwtCustomizer() {
        return new ReactiveJwtCustomizer() {

        };
    }


}
